import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
# 修复matplotlib中文显示问题
import matplotlib
matplotlib.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans', 'Arial Unicode MS', 'sans-serif']
matplotlib.rcParams['axes.unicode_minus'] = False
plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans', 'Arial Unicode MS', 'sans-serif']
plt.rcParams['axes.unicode_minus'] = False

import torch
from torch.utils.data import Dataset, DataLoader
from sklearn.metrics import accuracy_score, confusion_matrix
from sklearn.preprocessing import StandardScaler

# 导入所有模型
from trainCode.RAE_GCN_Transformer_BSM1 import RAE_GCN_Transformer, BSM1Dataset
from trainCode.SVM_BSM1 import main as svm_main
from trainCode.LDA_BSM1 import main as lda_main
from trainCode.Transformer_BSM1 import Transformer
from trainCode.GCN_BSM1 import GCN
from trainCode.GCN_Transformer_BSM1 import GCN_Transformer
from trainCode.RAE_Transformer_BSM1 import RAE_Transformer

# =========================================================================
# 模型内部名称到显示名称的映射
# =========================================================================
DISPLAY_NAME_MAP = {
    "LDA": "LDA",
    "SVM": "SVM",
    "RAE_Transformer": "RAE", # 对应图片中的 'RAE'
    "GCN": "GCN",
    "Transformer": "Transformer",
    "RAE_Transformer_v2": "RAE-Transformer", # 对应图片中的 'RAE-Transformer'
    "GCN_Transformer": "GCN-Transformer",
    "RAE_GCN_Transformer": "RAE-GCN-Transformer"
}

# =========================================================================
# 硬编码的准确率数据，来源于用户提供的表格图片 (RAE_GCN_Transformer 除外)
# 注意：RAE对应代码中的RAE_Transformer, RAE-Transformer对应RAE_Transformer_v2, RAE-GCN-Transformer是实际评估
# =========================================================================
HARDCODED_ACCURACIES = {
    "LDA": {
        "overall": 0.736,
        "class_wise": [0.595, 0.836, 0.800, 0.835, 0.895, 0.775, 0.825, 0.400, 0.565, 0.665, 0.905]
    },
    "SVM": {
        "overall": 0.728,
        "class_wise": [0.650, 0.865, 0.820, 0.820, 0.915, 0.755, 0.705, 0.260, 0.555, 0.750, 0.910]
    },
    "RAE_Transformer": { # 对应图片中的 'RAE'
        "overall": 0.812,
        "class_wise": [0.700, 0.905, 0.835, 0.845, 0.930, 0.835, 0.855, 0.430, 0.830, 0.825, 0.945]
    },
    "GCN": {
        "overall": 0.837,
        "class_wise": [0.775, 0.930, 0.865, 0.915, 0.945, 0.830, 0.860, 0.520, 0.735, 0.870, 0.965]
    },
    "Transformer": {
        "overall": 0.850,
        "class_wise": [0.880, 0.920, 0.865, 0.890, 0.925, 0.831, 0.835, 0.560, 0.825, 0.890, 0.930]
    },
    "RAE_Transformer_v2": { # 对应图片中的 'RAE-Transformer'
        "overall": 0.897,
        "class_wise": [0.901, 0.946, 0.906, 0.936, 0.956, 0.861, 0.876, 0.791, 0.841, 0.896, 0.956]
    },
    "GCN_Transformer": { # 对应图片中的 'GCN-Transformer'
        "overall": 0.907,
        "class_wise": [0.902, 0.977, 0.917, 0.928, 0.967, 0.871, 0.862, 0.807, 0.837, 0.937, 0.977]
    }
}

def load_model_results(model_name):
    """加载模型的评估结果"""
    try:
        acc_val = np.load(f'model/acc_val_{model_name}_BSM1.npy')
        return acc_val
    except FileNotFoundError:
        print(f"警告: 未找到模型 {model_name} 的验证准确率文件，将无法绘制学习曲线。")
        return None

def evaluate_model(model, model_name, device):
    """评估模型在测试集上的性能
    对于 RAE_GCN_Transformer，执行实际评估；对于其他模型，使用硬编码数据。
    """
    if model_name == "RAE_GCN_Transformer":
        print(f"\n正在评估 {model_name} 模型...")
        # 检查模型参数
        if hasattr(model, 'embed_dim') and hasattr(model, 'num_heads'):
            embed_dim = getattr(model, 'embed_dim')
            num_heads = getattr(model, 'num_heads')
            if embed_dim % num_heads != 0:
                print(f"警告: 模型 {model_name} 的 embed_dim({embed_dim}) 不能被 num_heads({num_heads}) 整除")
                print("这可能导致AssertionError: embed_dim must be divisible by num_heads错误")
                return None, None
            else:
                print(f"模型参数检查：embed_dim({embed_dim}) 可以被 num_heads({num_heads}) 整除")
        
        # 加载测试数据
        test_data = BSM1Dataset('BSM1_WWTP_data/test_data.csv', seq_len=10)
        test_loader = DataLoader(test_data, batch_size=512, shuffle=False)
        
        # 加载模型权重
        try:
            model.load_state_dict(torch.load(f'model/model_{model_name}_BSM1.pth', map_location=device))
            model.eval()
        except FileNotFoundError:
            print(f"错误: 未找到模型 {model_name} 的权重文件")
            return None, None
        except RuntimeError as e:
            print(f"错误: 无法加载模型 {model_name} 的权重，可能是模型结构与保存的权重不匹配")
            print(f"详细错误: {str(e)}")
            print(f"跳过模型 {model_name} 的评估")
            return None, None
        
        # 在测试集上评估
        try:
            y_true = []
            y_pred = []
            
            with torch.no_grad():
                for x, y in test_loader:
                    x = x.to(device)
                    y_true.extend(y.numpy())
                    outputs = model(x)
                    _, predicted = torch.max(outputs.data, 1)
                    y_pred.extend(predicted.cpu().numpy())
        except Exception as e:
            print(f"错误: 模型 {model_name} 在推理过程中出错")
            print(f"详细错误: {str(e)}")
            return None, None
        
        # 计算准确率和混淆矩阵
        accuracy = accuracy_score(y_true, y_pred)
        cm = confusion_matrix(y_true, y_pred)
        class_accuracies = cm.diagonal() / cm.sum(axis=1)
        
        # 打印总体准确率
        print(f"{model_name} 测试集准确率: {accuracy:.4f}")
        
        # 打印每个类别的准确率
        print(f"\n{model_name} 在各故障类别上的准确率:")
        print("-" * 50)
        print("类别\t准确率")
        for i, acc in enumerate(class_accuracies):
            print(f"{i}\t{acc:.4f}")
        print(f"平均\t{accuracy:.4f}")
        print("-" * 50)
        
        return accuracy, class_accuracies
    else:
        # 对于其他模型，使用硬编码数据
        print(f"\n正在使用硬编码数据评估 {model_name} 模型...")
        if model_name in HARDCODED_ACCURACIES:
            overall_acc = HARDCODED_ACCURACIES[model_name]["overall"]
            class_acc_list = HARDCODED_ACCURACIES[model_name]["class_wise"]
            
            # 打印总体准确率
            print(f"{model_name} 测试集准确率 (硬编码): {overall_acc:.4f}")
            
            # 打印每个类别的准确率
            print(f"\n{model_name} 在各故障类别上的准确率 (硬编码):")
            print("-" * 50)
            print("类别\t准确率")
            for i, acc in enumerate(class_acc_list):
                print(f"{i}\t{acc:.4f}")
            print(f"平均\t{overall_acc:.4f}")
            print("-" * 50)
            
            return overall_acc, np.array(class_acc_list)
        else:
            print(f"错误: 模型 {model_name} 未在硬编码数据中找到。")
            return None, None

def compare_models():
    """比较所有模型的性能"""
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"使用设备: {device}")
    
    # 获取数据集信息 (仍然需要num_classes用于表格生成)
    train_data = BSM1Dataset('BSM1_WWTP_data/train_data.csv', seq_len=10)
    input_dim = train_data.data.shape[1] # 保持此行，用于 RAE_GCN_Transformer 初始化
    num_classes = len(np.unique(train_data.labels))
    seq_len = 10 # 保持此行，用于 RAE_GCN_Transformer 初始化

    results = {}
    class_results = {}
    
    # 只实例化需要实际评估的模型
    print("正在实例化 RAE_GCN_Transformer 模型...")
    rae_gcn_transformer_model = RAE_GCN_Transformer(input_dim=input_dim, seq_len=seq_len, num_classes=num_classes).to(device)

    # --- 1. 处理 SVM 和 LDA （它们有自己的main函数来生成结果文件）---
    print("\n运行SVM模型评估...")
    svm_main() # 假设会生成 acc_test_SVM_BSM1.npy 和 class_acc_SVM_BSM1.npy
    print("\n运行LDA模型评估...")
    lda_main() # 假设会生成 acc_test_LDA_BSM1.npy 和 class_acc_LDA_BSM1.npy
    
    # 从 HARDCODED_ACCURACIES 加载 SVM 和 LDA 的结果 (确保与硬编码匹配)
    for model_name in ["LDA", "SVM"]:
        if model_name in HARDCODED_ACCURACIES:
            results[model_name] = HARDCODED_ACCURACIES[model_name]["overall"]
            class_results[model_name] = np.array(HARDCODED_ACCURACIES[model_name]["class_wise"])
            print(f"已加载 {model_name} 的硬编码准确率: {results[model_name]:.4f}")
        else:
            print(f"错误: 硬编码数据中缺少模型 {model_name} 的信息。")

    # --- 2. 处理其他硬编码的深度学习模型 ---
    # 除了 RAE_GCN_Transformer，其他深度学习模型从硬编码数据中获取
    hardcoded_dl_models = [name for name in HARDCODED_ACCURACIES.keys() if name not in ["LDA", "SVM"]]
    
    for model_name in hardcoded_dl_models:
        if model_name in HARDCODED_ACCURACIES:
            results[model_name] = HARDCODED_ACCURACIES[model_name]["overall"]
            class_results[model_name] = np.array(HARDCODED_ACCURACIES[model_name]["class_wise"])
            print(f"已加载 {model_name} 的硬编码准确率: {results[model_name]:.4f}")
        else:
            print(f"警告: 硬编码数据中缺少模型 {model_name} 的信息。")
    
    # --- 3. 实际评估 RAE_GCN_Transformer 模型 ---
    print("\n开始评估 RAE_GCN_Transformer 模型...")
    rae_gcn_transformer_acc, rae_gcn_transformer_class_acc = evaluate_model(rae_gcn_transformer_model, 'RAE_GCN_Transformer', device)
    if rae_gcn_transformer_acc is not None:
        results['RAE_GCN_Transformer'] = rae_gcn_transformer_acc
        class_results['RAE_GCN_Transformer'] = rae_gcn_transformer_class_acc
    else:
        print("RAE_GCN_Transformer 评估失败，跳过该模型")

    # 创建结果目录
    os.makedirs('results', exist_ok=True)
    
    # 保存结果
    # 定义模型显示顺序（RAE_GCN_Transformer放到最后, 与图片顺序一致）
    model_display_order = ["LDA", "SVM", "RAE_Transformer", "GCN", "Transformer", "RAE_Transformer_v2", "GCN_Transformer", "RAE_GCN_Transformer"]
    
    # 过滤掉results中不在model_display_order的键，并按顺序排序
    ordered_results = {k: results[k] for k in model_display_order if k in results}

    result_df = pd.DataFrame({
        'Model': [DISPLAY_NAME_MAP.get(k, k) for k in ordered_results.keys()],
        'Accuracy': list(ordered_results.values())
    })
    result_df.to_csv('results/model_comparison.csv', index=False)
    
    # 打印每个故障类别的准确率
    print("\n各模型在每个故障类别上的准确率:")
    print("-" * 80)
    header = "模型名称"
    for i in range(num_classes):
        header += f"\t{i}"
    header += "\t平均"
    print(header)
    
    for model_name in model_display_order: # 使用model_display_order来控制打印顺序
        if model_name in class_results:
            display_name = DISPLAY_NAME_MAP.get(model_name, model_name)
            class_acc = class_results[model_name]
            row = f"{display_name:<20}" # 调整对齐
            for i in range(num_classes):
                row += f"\t{class_acc[i]:.4f}"
            row += f"\t{results[model_name]:.4f}"
            print(row)
    print("-" * 80)
    
    # 生成表格4-4：不同模型对各故障及正常数据的分类效果
    generate_classification_table(results, class_results, num_classes)
    
    # 绘制模型比较图
    plt.figure(figsize=(12, 6))
    bars = plt.bar(result_df['Model'], result_df['Accuracy'])
    plt.xlabel('模型')
    plt.ylabel('测试集准确率')
    plt.title('不同模型在BSM1数据集上的性能比较')
    plt.xticks(rotation=45)
    plt.ylim(0, 1.0)
    
    # 在柱状图上添加数值标签
    for bar in bars:
        height = bar.get_height()
        plt.text(bar.get_x() + bar.get_width()/2., height + 0.01,
                f'{height:.4f}', ha='center', va='bottom')
    
    plt.tight_layout()
    plt.savefig('results/model_comparison.png', dpi=300)
    
    # 绘制每个模型在各类别上的准确率
    plt.figure(figsize=(15, 10))
    x = np.arange(num_classes)
    width = 0.1
    offset = -width * (len(class_results) - 1) / 2
    
    for i, (name, class_acc) in enumerate(class_results.items()):
        display_name = DISPLAY_NAME_MAP.get(name, name)
        plt.bar(x + offset + i*width, class_acc, width, label=display_name)
    
    plt.xlabel('类别')
    plt.ylabel('准确率')
    plt.title('不同模型在各类别上的准确率比较')
    plt.xticks(x)
    plt.ylim(0, 1.0)
    plt.legend()
    plt.tight_layout()
    plt.savefig('results/class_accuracy_comparison.png', dpi=300)
    
    # 绘制验证集上的学习曲线
    plt.figure(figsize=(12, 6))
    # 遍历所有模型名称来加载学习曲线数据
    for name in model_display_order: # 使用 model_display_order 来确保顺序
        acc_val = load_model_results(name) # 尝试加载真实的验证准确率，如果文件不存在会打印警告
        if acc_val is not None:
            display_name = DISPLAY_NAME_MAP.get(name, name)
            plt.plot(acc_val, label=display_name)
    
    plt.xlabel('训练轮次')
    plt.ylabel('验证集准确率')
    plt.title('不同模型在验证集上的学习曲线')
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    plt.savefig('results/learning_curves.png', dpi=300)
    
    print("\n模型比较完成，结果已保存到 'results' 目录")

def generate_classification_table(results, class_results, num_classes):
    """生成表格4-4：不同模型对各故障及正常数据的分类效果"""
    print("\n正在生成表格4-4：不同模型对各故障及正常数据的分类效果...")
    
    # 创建DataFrame用于保存表格数据
    # 设置列名：类别0-10和平均
    columns = [f'{i}' for i in range(num_classes)] + ['平均']
    
    # 创建空的DataFrame
    table_df = pd.DataFrame(columns=columns)
    
    # 定义模型显示顺序（与图片一致）
    model_order = ["LDA", "SVM", "RAE_Transformer", "GCN", "Transformer", "RAE_Transformer_v2", "GCN_Transformer", "RAE_GCN_Transformer"]
    
    # 按照指定顺序填充数据（以百分比形式）
    for model_name in model_order:
        if model_name in class_results:
            class_acc = class_results[model_name]
            row_data = {}
            # 添加各类别准确率（百分比形式）
            for i in range(num_classes):
                row_data[f'{i}'] = f'{class_acc[i]*100:.1f}%'
            # 添加平均准确率
            row_data['平均'] = f'{results[model_name]*100:.1f}%'
            # 将行数据添加到DataFrame，使用显示名称作为索引
            display_name = DISPLAY_NAME_MAP.get(model_name, model_name)
            table_df.loc[display_name] = row_data
    
    # 重命名索引为"模型"
    table_df.index.name = '模型'
    
    # 保存为CSV文件
    try:
        # 确保results目录存在
        os.makedirs('results', exist_ok=True)
        
        csv_path = 'results/table4-4_classification_results.csv'
        table_df.to_csv(csv_path, encoding='utf-8-sig')  # 使用utf-8-sig编码以支持中文
        print(f"表格已保存到: {csv_path}")
    except PermissionError:
        print(f"警告: 无法保存表格到文件，可能是文件被占用或没有足够权限")
    except Exception as e:
        print(f"警告: 保存表格时出错: {str(e)}")
    
    # 打印表格内容
    print("\n表格4-4 不同模型对各故障及正常数据的分类效果")
    print("Table 4-4 Classification results of different models for normal data and faults")
    print(table_df.to_string())
    
    return table_df

# 修改SVM和LDA的main函数，使其可以被导入而不自动执行
def patch_main_functions():
    with open('trainCode/SVM_BSM1.py', 'r', encoding='utf-8') as f:
        svm_content = f.read()
    
    svm_content = svm_content.replace('def main():', 'def main(run_plots=True):')
    svm_content = svm_content.replace('plt.show()', 'if run_plots:\n        plt.show()')
    
    with open('trainCode/SVM_BSM1.py', 'w', encoding='utf-8') as f:
        f.write(svm_content)
    
    with open('trainCode/LDA_BSM1.py', 'r', encoding='utf-8') as f:
        lda_content = f.read()
    
    lda_content = lda_content.replace('def main():', 'def main(run_plots=True):')
    lda_content = lda_content.replace('plt.show()', 'if run_plots:\n        plt.show()')
    
    with open('trainCode/LDA_BSM1.py', 'w', encoding='utf-8') as f:
        f.write(lda_content)

if __name__ == '__main__':
    # 修补main函数
    patch_main_functions()
    
    # 比较所有模型
    compare_models()